聊聊Spring AI的RetrievalAugmentationAdvisor

序

本文主要研究一下Spring AI的RetrievalAugmentationAdvisor

BaseAdvisor

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/api/BaseAdvisor.java

public interface BaseAdvisor extends CallAroundAdvisor, StreamAroundAdvisor {

    Scheduler DEFAULT_SCHEDULER = Schedulers.boundedElastic();

    @Override
    default AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        Assert.notNull(advisedRequest, "advisedRequest cannot be null");
        Assert.notNull(chain, "chain cannot be null");

        AdvisedRequest processedAdvisedRequest = before(advisedRequest);
        AdvisedResponse advisedResponse = chain.nextAroundCall(processedAdvisedRequest);
        return after(advisedResponse);
    }

    @Override
    default Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        Assert.notNull(advisedRequest, "advisedRequest cannot be null");
        Assert.notNull(chain, "chain cannot be null");
        Assert.notNull(getScheduler(), "scheduler cannot be null");

        Flux<AdvisedResponse> advisedResponses = Mono.just(advisedRequest)
            .publishOn(getScheduler())
            .map(this::before)
            .flatMapMany(chain::nextAroundStream);

        return advisedResponses.map(ar -> {
            if (AdvisedResponseStreamUtils.onFinishReason().test(ar)) {
                ar = after(ar);
            }
            return ar;
        }).onErrorResume(error -> Flux.error(new IllegalStateException("Stream processing failed", error)));
    }

    @Override
    default String getName() {
        return this.getClass().getSimpleName();
    }

    /**
     * Logic to be executed before the rest of the advisor chain is called.
     */
    AdvisedRequest before(AdvisedRequest request);

    /**
     * Logic to be executed after the rest of the advisor chain is called.
     */
    AdvisedResponse after(AdvisedResponse advisedResponse);

    /**
     * Scheduler used for processing the advisor logic when streaming.
     */
    default Scheduler getScheduler() {
        return DEFAULT_SCHEDULER;
    }

}

BaseAdvisor继承了CallAroundAdvisor、StreamAroundAdvisor，它提供了aroundCall、aroundStream的default，在执行chain的next之前执行before，之后执行after方法

RetrievalAugmentationAdvisor

spring-ai-core/src/main/java/org/springframework/ai/chat/client/advisor/RetrievalAugmentationAdvisor.java

public final class RetrievalAugmentationAdvisor implements BaseAdvisor {

    public static final String DOCUMENT_CONTEXT = "rag_document_context";

    private final List<QueryTransformer> queryTransformers;

    @Nullable
    private final QueryExpander queryExpander;

    private final DocumentRetriever documentRetriever;

    private final DocumentJoiner documentJoiner;

    private final QueryAugmenter queryAugmenter;

    private final TaskExecutor taskExecutor;

    private final Scheduler scheduler;

    private final int order;

    public RetrievalAugmentationAdvisor(@Nullable List<QueryTransformer> queryTransformers,
            @Nullable QueryExpander queryExpander, DocumentRetriever documentRetriever,
            @Nullable DocumentJoiner documentJoiner, @Nullable QueryAugmenter queryAugmenter,
            @Nullable TaskExecutor taskExecutor, @Nullable Scheduler scheduler, @Nullable Integer order) {
        Assert.notNull(documentRetriever, "documentRetriever cannot be null");
        Assert.noNullElements(queryTransformers, "queryTransformers cannot contain null elements");
        this.queryTransformers = queryTransformers != null ? queryTransformers : List.of();
        this.queryExpander = queryExpander;
        this.documentRetriever = documentRetriever;
        this.documentJoiner = documentJoiner != null ? documentJoiner : new ConcatenationDocumentJoiner();
        this.queryAugmenter = queryAugmenter != null ? queryAugmenter : ContextualQueryAugmenter.builder().build();
        this.taskExecutor = taskExecutor != null ? taskExecutor : buildDefaultTaskExecutor();
        this.scheduler = scheduler != null ? scheduler : BaseAdvisor.DEFAULT_SCHEDULER;
        this.order = order != null ? order : 0;
    }

    public static Builder builder() {
        return new Builder();
    }

    @Override
    public AdvisedRequest before(AdvisedRequest request) {
        Map<String, Object> context = new HashMap<>(request.adviseContext());

        // 0. Create a query from the user text, parameters, and conversation history.
        Query originalQuery = Query.builder()
            .text(new PromptTemplate(request.userText(), request.userParams()).render())
            .history(request.messages())
            .context(context)
            .build();

        // 1. Transform original user query based on a chain of query transformers.
        Query transformedQuery = originalQuery;
        for (var queryTransformer : this.queryTransformers) {
            transformedQuery = queryTransformer.apply(transformedQuery);
        }

        // 2. Expand query into one or multiple queries.
        List<Query> expandedQueries = this.queryExpander != null ? this.queryExpander.expand(transformedQuery)
                : List.of(transformedQuery);

        // 3. Get similar documents for each query.
        Map<Query, List<List<Document>>> documentsForQuery = expandedQueries.stream()
            .map(query -> CompletableFuture.supplyAsync(() -> getDocumentsForQuery(query), this.taskExecutor))
            .toList()
            .stream()
            .map(CompletableFuture::join)
            .collect(Collectors.toMap(Map.Entry::getKey, entry -> List.of(entry.getValue())));

        // 4. Combine documents retrieved based on multiple queries and from multiple data
        // sources.
        List<Document> documents = this.documentJoiner.join(documentsForQuery);
        context.put(DOCUMENT_CONTEXT, documents);

        // 5. Augment user query with the document contextual data.
        Query augmentedQuery = this.queryAugmenter.augment(originalQuery, documents);

        // 6. Update advised request with augmented prompt.
        return AdvisedRequest.from(request).userText(augmentedQuery.text()).adviseContext(context).build();
    }

    /**
     * Processes a single query by routing it to document retrievers and collecting
     * documents.
     */
    private Map.Entry<Query, List<Document>> getDocumentsForQuery(Query query) {
        List<Document> documents = this.documentRetriever.retrieve(query);
        return Map.entry(query, documents);
    }

    @Override
    public AdvisedResponse after(AdvisedResponse advisedResponse) {
        ChatResponse.Builder chatResponseBuilder;
        if (advisedResponse.response() == null) {
            chatResponseBuilder = ChatResponse.builder();
        }
        else {
            chatResponseBuilder = ChatResponse.builder().from(advisedResponse.response());
        }
        chatResponseBuilder.metadata(DOCUMENT_CONTEXT, advisedResponse.adviseContext().get(DOCUMENT_CONTEXT));
        return new AdvisedResponse(chatResponseBuilder.build(), advisedResponse.adviseContext());
    }

    @Override
    public Scheduler getScheduler() {
        return this.scheduler;
    }

    @Override
    public int getOrder() {
        return this.order;
    }

    private static TaskExecutor buildDefaultTaskExecutor() {
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setThreadNamePrefix("ai-advisor-");
        taskExecutor.setCorePoolSize(4);
        taskExecutor.setMaxPoolSize(16);
        taskExecutor.setTaskDecorator(new ContextPropagatingTaskDecorator());
        taskExecutor.initialize();
        return taskExecutor;
    }

    //......

}    

RetrievalAugmentationAdvisor实现了BaseAdvisor接口，其before方法先通过queryTransformers来转换原始query，之后使用queryExpander来扩展query，接着通过documentRetriever来对扩展query检索相关文档，然后通过documentJoiner.join来把检索结果结合起来，作为rag_document_context放到context中，最后通过queryAugmenter.augment(originalQuery, documents)来生成最后的augmentedQuery；其after方法将before存储下来的key为rag_document_context的检索到的文档作为metadata附加到response中。

示例

    @Test
    void ragWithRewrite() {
        String question = "Where are the main characters going?";

        RetrievalAugmentationAdvisor ragAdvisor = RetrievalAugmentationAdvisor.builder()
            .queryTransformers(RewriteQueryTransformer.builder()
                .chatClientBuilder(ChatClient.builder(this.openAiChatModel))
                .targetSearchSystem("vector store")
                .build())
            .documentRetriever(VectorStoreDocumentRetriever.builder().vectorStore(this.pgVectorStore).build())
            .build();

        ChatResponse chatResponse = ChatClient.builder(this.openAiChatModel)
            .build()
            .prompt()
            .user(question)
            .advisors(ragAdvisor)
            .call()
            .chatResponse();

        assertThat(chatResponse).isNotNull();

        String response = chatResponse.getResult().getOutput().getText();
        System.out.println(response);
        assertThat(response).containsIgnoringCase("Loch of the Stars");

        evaluateRelevancy(question, chatResponse);
    }

这里使用了RewriteQueryTransformer、VectorStoreDocumentRetriever

小结

Spring AI定义了BaseAdvisor，它继承了CallAroundAdvisor、StreamAroundAdvisor，它提供了aroundCall、aroundStream的default，在执行chain的next之前执行before，之后执行after方法。RetrievalAugmentationAdvisor实现了BaseAdvisor接口，其before方法先通过queryTransformers来转换原始query，之后使用queryExpander来扩展query，接着通过documentRetriever来对扩展query检索相关文档，然后通过documentJoiner.join来把检索结果结合起来，作为rag_document_context放到context中，最后通过queryAugmenter.augment(originalQuery, documents)来生成最后的augmentedQuery；其after方法将before存储下来的key为rag_document_context的检索到的文档作为metadata附加到response中。

doc

• Advisors API