深入探索Spring AI：原始碼分析流式回答

努力的小雨發表於2024-10-11
SpringAI原始碼

在上一章節中，我們深入分析了Spring AI的阻塞式請求與響應機制，並探討了如何增強其記憶能力。今天，我們將重點講解流式響應的概念與實現。畢竟，AI的流式回答功能與其互動體驗密切相關，是提升使用者滿意度的重要組成部分。

基本用法

基本用法非常簡單，只需增加一個 stream 方法即可實現所需功能。接下來，我們將透過程式碼示例來展示這一過程，幫助您更清晰地理解如何在實際應用中進行操作。請看以下程式碼：

@GetMapping(value = "/ai-stream",produces = MediaType.APPLICATION_OCTET_STREAM_VALUE + ";charset=UTF-8")
Flux<String> generationByStream(@RequestParam("userInput") String userInput) {
    Flux<String> output = chatClient.prompt()
            .user(userInput)
            .stream()
            .content();
    return output;
}

在我們增加 stream 方法之後，返回的物件型別將不再是原來的阻塞式 CallResponseSpec，而是轉換為非阻塞的 StreamResponseSpec。與此同時，返回的資料型別也由之前的 String 變更為 Flux。

在深入探討其具體應用之前，首先讓我來介紹一下 Flux 的概念與特性。

Spring WebFlux的處理器實現

首先，在 WebFlux 中，處理器已經實現了非阻塞式的功能。這意味著，只要我們的程式碼返回一個 Flux 物件，就能輕鬆實現響應功能。透過這種方式，應用程式能夠高效地處理併發請求，而不會因阻塞操作而影響整體效能。

    @Override
    public Mono<Void> handle(ServerWebExchange exchange) {
        if (this.handlerMappings == null) {
            return createNotFoundError();
        }
        if (CorsUtils.isPreFlightRequest(exchange.getRequest())) {
            return handlePreFlight(exchange);
        }
        return Flux.fromIterable(this.handlerMappings)
                .concatMap(mapping -> mapping.getHandler(exchange))
                .next()
                .switchIfEmpty(createNotFoundError())
                .onErrorResume(ex -> handleResultMono(exchange, Mono.error(ex)))
                .flatMap(handler -> handleRequestWith(exchange, handler));
    }

這裡簡單介紹一下 Spring WebFlux，雖然這不是我們的重點，但瞭解其基本概念還是很有幫助的。Spring WebFlux 是 Spring 框架的一部分，專為構建反應式應用而設計。它支援非同步和非阻塞的程式設計模型，使得處理高併發請求變得更加高效。以下是 WebFlux 的幾個關鍵特性：

  1. 反應式程式設計：WebFlux 基於反應式程式設計模型，使用 Mono 和 Flux 型別來處理資料流。Mono 表示零或一個元素，而 Flux 則表示零個或多個元素。這種模型使得我們可以輕鬆處理非同步資料流，從而提高程式碼的可讀性和可維護性。
  2. 非阻塞 I/O：WebFlux 透過非阻塞的 I/O 操作（如 Netty 或 Servlet 3.1+ 容器）來實現高效的資源利用。與傳統的阻塞 I/O 不同，WebFlux 在等待響應時能夠釋放執行緒，這樣一來，就可以顯著提高應用的併發能力，支援更多的同時請求而不增加執行緒開銷。

瞭解這些特性將為後續的非阻塞式響應設計奠定基礎，幫助我們更好地利用 WebFlux 的能力來提升應用效能。

原始碼分析

現在我們來詳細看看我們的 content 是如何操作的。接下來的程式碼示例將展示具體的實現方式，幫助我們理解在 WebFlux 中如何處理資料流和響應:

public Flux<String> content() {
    return doGetFluxChatResponse(this.request).map(r -> {
        if (r.getResult() == null || r.getResult().getOutput() == null
                || r.getResult().getOutput().getContent() == null) {
            return "";
        }
        return r.getResult().getOutput().getContent();
    }).filter(StringUtils::hasLength);
}

這裡的實現相對簡單，主要是傳入了一個函式。接下來，我們將深入分析 doGetFluxChatResponse 的程式碼實現，以便更好地理解其具體邏輯和運作方式：

private Flux<ChatResponse> doGetFluxChatResponse2(DefaultChatClientRequestSpec inputRequest) {
//此處省略重複程式碼
    var fluxChatResponse = this.chatModel.stream(prompt);
//此處省略重複程式碼
    return advisedResponse;
}

這裡的程式碼邏輯與阻塞回答基本相同，唯一的不同之處在於它呼叫了 chatModel.stream(prompt) 方法。接下來，我們將深入探討 chatModel.stream(prompt) 方法的具體實現和其背後的設計思路：

public Flux<ChatResponse> stream(Prompt prompt) {
        return Flux.deferContextual(contextView -> {
        //此處省略重複程式碼
            Flux<OpenAiApi.ChatCompletionChunk> completionChunks = this.openAiApi.chatCompletionStream(request,
                    getAdditionalHttpHeaders(prompt));
//此處省略重複程式碼
            Flux<ChatResponse> chatResponse = completionChunks.map(this::chunkToChatCompletion)
                .switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> {
//此處省略重複程式碼
                        return new ChatResponse(generations, from(chatCompletion2, null));
                    }
                }));
//此處省略重複程式碼
            return new MessageAggregator().aggregate(flux, observationContext::setResponse);

        });
    }

同樣的邏輯在這裡就不再贅述，我們將重點關注其中的區別。在這一部分，我們使用了 chatCompletionStream，而且與之前不同的是，這裡不再使用 retryTemplate，而是引入了 webClient，這是一個能夠接收事件流的工具類。

public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest,
        MultiValueMap<String, String> additionalHttpHeader) {

    Assert.notNull(chatRequest, "The request body can not be null.");
    Assert.isTrue(chatRequest.stream(), "Request must set the stream property to true.");

    AtomicBoolean isInsideTool = new AtomicBoolean(false);

    return this.webClient.post()
        .uri(this.completionsPath)
        .headers(headers -> headers.addAll(additionalHttpHeader))
        .body(Mono.just(chatRequest), ChatCompletionRequest.class)
        .retrieve()
        .bodyToFlux(String.class)
        // cancels the flux stream after the "[DONE]" is received.
        .takeUntil(SSE_DONE_PREDICATE)
        // filters out the "[DONE]" message.
        .filter(SSE_DONE_PREDICATE.negate())
        .map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class))
//此處省略一堆程式碼

這段程式碼的主要目的是透過 webClient 向指定路徑發起一個 POST 請求，同時設定合適的請求頭和請求體。在獲取響應資料時，使用了事件流的方式（透過 bodyToFlux 方法）來接收響應內容，並對資料進行過濾和轉換，最終將其轉化為 ChatCompletionChunk 物件。

儘管其餘的業務邏輯與之前相似，但有一點顯著的區別，即整個流程的返回型別以及與 OpenAI API 的呼叫方式都是非阻塞式的。

總結

在當今的數字時代，流式響應機制不僅提升了系統的效能，還在使用者體驗上扮演了關鍵角色。透過引入 Flux 型別，Spring WebFlux 的設計理念使得應用能夠以非阻塞的方式處理併發請求，從而有效利用資源並減少響應延遲。

我們終於全面講解了Spring AI的基本操作，包括阻塞式回答、流式回答以及記憶增強功能。這些內容為我們深入理解其工作機制奠定了基礎。接下來，我們將繼續深入探索原始碼，重點分析回撥函式、實體類對映等重要功能。

這將幫助我們更好地理解Spring AI的內部運作原理，併為進一步的最佳化和定製化提供指導。

我是努力的小雨，一名 Java 服務端碼農，潛心研究著 AI 技術的奧秘。我熱愛技術交流與分享，對開源社群充滿熱情。同時也是一位騰訊雲創作之星、阿里雲專家博主、華為云云享專家、掘金優秀作者。

💡 我將不吝分享我在技術道路上的個人探索與經驗，希望能為你的學習與成長帶來一些啟發與幫助。

🌟 歡迎關注努力的小雨！🌟

相關文章