com.microsoft.semantickernel.data.VectorStoreTextSearch Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of semantickernel-api Show documentation
Defines the public interface for the Semantic Kernel
There is a newer version: 1.4.2
// Copyright (c) Microsoft. All rights reserved.
package com.microsoft.semantickernel.data;

import com.microsoft.semantickernel.data.textsearch.DefaultTextSearchResultMapper;
import com.microsoft.semantickernel.data.textsearch.DefaultTextSearchStringMapper;
import com.microsoft.semantickernel.data.textsearch.KernelSearchResults;
import com.microsoft.semantickernel.data.textsearch.TextSearch;
import com.microsoft.semantickernel.data.textsearch.TextSearchOptions;
import com.microsoft.semantickernel.data.textsearch.TextSearchResult;
import com.microsoft.semantickernel.data.textsearch.TextSearchResultMapper;
import com.microsoft.semantickernel.data.textsearch.TextSearchStringMapper;
import com.microsoft.semantickernel.data.vectorsearch.VectorSearchFilter;
import com.microsoft.semantickernel.data.vectorsearch.VectorSearchResults;
import com.microsoft.semantickernel.data.vectorsearch.VectorizedSearch;
import com.microsoft.semantickernel.data.vectorstorage.options.VectorSearchOptions;
import com.microsoft.semantickernel.exceptions.SKException;
import com.microsoft.semantickernel.services.textembedding.TextEmbeddingGenerationService;
import reactor.core.publisher.Mono;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.stream.Collectors;

/**
 * A text search implementation that uses a vector record collection to perform the search.
 *
 * @param  The record type.
 */
public class VectorStoreTextSearch implements TextSearch {

    private final VectorizedSearch vectorizedSearch;
    private final TextEmbeddingGenerationService textEmbeddingGenerationService;
    private final TextSearchStringMapper stringMapper;
    private final TextSearchResultMapper resultMapper;
    private final VectorStoreTextSearchOptions options;

    /**
     * Create a new instance of VectorStoreTextSearch with {@link VectorizedSearch} to perform vectorized search and
     * {@link TextEmbeddingGenerationService} to generate text embeddings.
     *
     * @param vectorizedSearch               The vectorized search. Usually a vector record collection.
     * @param textEmbeddingGenerationService The text embedding generation service.
     * @param stringMapper                   The string mapper.
     * @param resultMapper                   The result mapper.
     * @param options                        The options.
     */
    public VectorStoreTextSearch(
        @Nonnull VectorizedSearch vectorizedSearch,
        @Nonnull TextEmbeddingGenerationService textEmbeddingGenerationService,
        @Nullable TextSearchStringMapper stringMapper,
        @Nullable TextSearchResultMapper resultMapper,
        @Nullable VectorStoreTextSearchOptions options) {
        this.vectorizedSearch = vectorizedSearch;
        this.textEmbeddingGenerationService = textEmbeddingGenerationService;
        this.stringMapper = stringMapper == null ? new DefaultTextSearchStringMapper()
            : stringMapper;
        this.resultMapper = resultMapper == null ? new DefaultTextSearchResultMapper()
            : resultMapper;
        this.options = options == null ? new VectorStoreTextSearchOptions() : options;
    }

    private Mono> executeSearchAsync(String query,
        TextSearchOptions options) {
        if (options == null) {
            options = TextSearchOptions.createDefault();
        }

        VectorSearchOptions vectorSearchOptions = VectorSearchOptions.builder()
            .withVectorSearchFilter(options.getFilter() != null
                ? new VectorSearchFilter(options.getFilter().getFilterClauses())
                : null)
            .withTop(options.getTop())
            .withSkip(options.getSkip())
            .withIncludeTotalCount(options.isIncludeTotalCount())
            .build();

        return textEmbeddingGenerationService.generateEmbeddingAsync(query)
            .flatMap(embedding -> vectorizedSearch.searchAsync(embedding.getVector(),
                vectorSearchOptions));
    }

    /**
     * Perform a search for content related to the specified query and return String values representing the search results.
     *
     * @param query   The text to search for.
     * @param options The search options.
     * @return The search results.
     */
    @Override
    public Mono> searchAsync(String query, TextSearchOptions options) {
        return executeSearchAsync(query, options)
            .map(results -> new KernelSearchResults<>(
                results.getResults().stream()
                    .map(r -> stringMapper.fromResultToString(r.getRecord()))
                    .collect(Collectors.toList()),
                results.getTotalCount(),
                results.getMetadata()));
    }

    /**
     * Perform a search for content related to the specified query and return TextSearchResult values representing the search results.
     *
     * @param query   The text to search for.
     * @param options The search options.
     * @return The search results.
     */
    @Override
    public Mono> getTextSearchResultsAsync(String query,
        TextSearchOptions options) {
        return executeSearchAsync(query, options)
            .map(results -> new KernelSearchResults<>(
                results.getResults().stream()
                    .map(r -> resultMapper.fromResultToTextSearchResult(r.getRecord()))
                    .collect(Collectors.toList()),
                results.getTotalCount(),
                results.getMetadata()));
    }

    /**
     * Perform a search for content related to the specified query and return Object values representing the search results.
     *
     * @param query   The text to search for.
     * @param options The search options.
     * @return The search results.
     */
    @Override
    public Mono> getSearchResultsAsync(String query,
        TextSearchOptions options) {
        return executeSearchAsync(query, options)
            .map(results -> new KernelSearchResults<>(
                results.getResults().stream()
                    .map(r -> resultMapper.fromResultToTextSearchResult(r.getRecord()))
                    .collect(Collectors.toList()),
                results.getTotalCount(),
                results.getMetadata()));
    }

    /**
     * Create a new instance of {@link Builder}.
     *
     * @param  The record type.
     * @return The builder.
     */
    public static  Builder builder() {
        return new Builder<>();
    }

    /**
     * A builder for the {@link VectorStoreTextSearch} class.
     *
     * @param  The record type.
     */
    public static class Builder {
        @Nullable
        private VectorizedSearch vectorizedSearch;
        @Nullable
        private TextEmbeddingGenerationService textEmbeddingGenerationService;
        @Nullable
        private TextSearchStringMapper stringMapper;
        @Nullable
        private TextSearchResultMapper resultMapper;
        @Nullable
        private VectorStoreTextSearchOptions options;

        /**
         * Sets the vectorized search.
         *
         * @param vectorizedSearch The vectorized search.
         * @return The builder.
         */
        public Builder withVectorizedSearch(VectorizedSearch vectorizedSearch) {
            this.vectorizedSearch = vectorizedSearch;
            return this;
        }

        /**
         * Sets the text embedding generation service.
         *
         * @param textEmbeddingGenerationService The text embedding generation service.
         * @return The builder.
         */
        public Builder withTextEmbeddingGenerationService(
            TextEmbeddingGenerationService textEmbeddingGenerationService) {
            this.textEmbeddingGenerationService = textEmbeddingGenerationService;
            return this;
        }

        /**
         * Sets the string mapper.
         *
         * @param stringMapper The string mapper.
         * @return The builder.
         */
        public Builder withStringMapper(TextSearchStringMapper stringMapper) {
            this.stringMapper = stringMapper;
            return this;
        }

        /**
         * Sets the result mapper.
         *
         * @param resultMapper The result mapper.
         * @return The builder.
         */
        public Builder withResultMapper(TextSearchResultMapper resultMapper) {
            this.resultMapper = resultMapper;
            return this;
        }

        /**
         * Sets the options.
         *
         * @param options The options.
         * @return The builder.
         */
        public Builder withOptions(VectorStoreTextSearchOptions options) {
            this.options = options;
            return this;
        }

        /**
         * Builds the {@link VectorStoreTextSearch} instance.
         *
         * @return The {@link VectorStoreTextSearch} instance.
         */
        public VectorStoreTextSearch build() {
            if (vectorizedSearch == null) {
                throw new SKException("Vectorized search is required");
            }
            if (textEmbeddingGenerationService == null) {
                throw new SKException("Text embedding generation service is required");
            }

            return new VectorStoreTextSearch<>(vectorizedSearch, textEmbeddingGenerationService,
                stringMapper, resultMapper, options);
        }
    }
}