Streaming response từ LLM API với Java 21 và Spring Boot

1. Tổng quan

Một câu trả lời dài từ LLM có thể mất 10–30 giây mới sinh xong. Nếu bạn chờ đủ cả câu rồi mới trả về, người dùng sẽ nhìn màn hình trống suốt thời gian đó. Streaming cho phép nhận từng token ngay khi model sinh ra, nên chữ đầu tiên thường xuất hiện sau chưa tới một giây.

Bài này hướng dẫn bạn đọc stream từ một API tương thích OpenAI Chat Completions bằng java.net.http.HttpClient, sau đó bằng Spring WebClient. Tiếp theo, bạn sẽ đẩy stream về trình duyệt qua Server-Sent Events (SSE) và viết test cho cả luồng.

2. Định dạng stream của LLM API

Khi gửi request với "stream": true, API không trả về một JSON duy nhất. Thay vào đó, nó trả về Content-Type: text/event-stream, và mỗi sự kiện là một dòng data::

data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":""}}]}

data: {"id":"c1","choices":[{"index":0,"delta":{"content":"Xin"}}]}

data: {"id":"c1","choices":[{"index":0,"delta":{"content":" chào"}}]}

data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}

data: [DONE]

Có ba điểm bạn cần nhớ:

  • Phần text nằm ở choices[0].delta.content. Một số chunk không có field này, ví dụ chunk đầu tiên chỉ có role, hoặc chunk cuối chỉ có finish_reason.
  • Các sự kiện được ngăn cách bởi một dòng trống.
  • data: [DONE] báo hiệu stream đã kết thúc. Giá trị này không phải JSON, nên bạn không được đưa nó vào Jackson.

3. Tách logic parse ra một lớp riêng

Logic parse nên nằm trong một lớp độc lập để dễ test. Lớp này chỉ cần Jackson, vốn đã có sẵn trong mọi Spring Boot starter web:

package dev.nonla.chat;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.Optional;

public class ChatChunkParser {

    public static final String DONE = "[DONE]";
    private final ObjectMapper mapper = new ObjectMapper();

    /** Nhận nguyên một dòng SSE, ví dụ: data: {...} */
    public Optional<String> parseLine(String line) {
        if (!line.startsWith("data:")) {
            return Optional.empty();
        }
        return parsePayload(line.substring(5).trim());
    }

    /** Nhận phần payload sau "data:" */
    public Optional<String> parsePayload(String payload) {
        if (payload.isEmpty() || DONE.equals(payload)) {
            return Optional.empty();
        }
        try {
            JsonNode content = mapper.readTree(payload)
                    .path("choices").path(0).path("delta").path("content");
            return content.isTextual() && !content.asText().isEmpty()
                    ? Optional.of(content.asText())
                    : Optional.empty();
        } catch (JsonProcessingException e) {
            throw new IllegalStateException("Chunk không hợp lệ: " + payload, e);
        }
    }
}

Cách lớp này hoạt động:

  • parseLine bỏ qua dòng trống và các dòng như event: hoặc : keep-alive.
  • path(...) của Jackson không ném exception khi thiếu node. Nếu chunk không có content, kết quả là Optional.empty().
  • [DONE] được chặn trước khi gọi readTree, nên không gây lỗi parse.

4. Streaming với java.net.http.HttpClient

Nếu bạn chỉ cần một CLI hoặc một job đơn giản, HttpClient có sẵn trong JDK là đủ. BodyHandlers.ofLines() trả về một Stream<String> được đọc dần theo dữ liệu đến từ mạng, không phải đợi tải xong toàn bộ body:

package dev.nonla.chat;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpResponse.BodyHandlers;
import java.time.Duration;
import java.util.Optional;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class StreamingChatDemo {

    public static void main(String[] args) throws Exception {
        String baseUrl = System.getenv("LLM_BASE_URL");
        String apiKey = System.getenv("LLM_API_KEY");

        String body = """
            {"model": "gpt-4o-mini", "stream": true,
             "messages": [{"role": "user",
                           "content": "Giải thích quan hệ Order và Customer trong 3 câu"}]}
            """;

        HttpRequest request = HttpRequest.newBuilder(URI.create(baseUrl + "/v1/chat/completions"))
                .header("Authorization", "Bearer " + apiKey)
                .header("Content-Type", "application/json")
                .header("Accept", "text/event-stream")
                .timeout(Duration.ofSeconds(60))
                .POST(HttpRequest.BodyPublishers.ofString(body))
                .build();

        HttpClient client = HttpClient.newHttpClient();
        HttpResponse<Stream<String>> response = client.send(request, BodyHandlers.ofLines());

        if (response.statusCode() != 200) {
            String error = response.body().collect(Collectors.joining("\n"));
            throw new IllegalStateException("HTTP " + response.statusCode() + ": " + error);
        }

        ChatChunkParser parser = new ChatChunkParser();
        try (Stream<String> lines = response.body()) {
            lines.takeWhile(line -> !line.equals("data: [DONE]"))
                 .map(parser::parseLine)
                 .flatMap(Optional::stream)
                 .forEach(token -> {
                     System.out.print(token);
                     System.out.flush();
                 });
        }
        System.out.println();
    }
}

Giải thích các dòng chính:

  • Status code phải được kiểm tra trước. Khi API trả lỗi (401, 429…), body là JSON lỗi chứ không phải SSE.
  • takeWhile dừng stream khi gặp [DONE]. Khối try-with-resources đóng stream, và kết nối được giải phóng theo.
  • System.out.flush() giúp chữ hiện ra ngay. Thiếu dòng này, terminal có thể gom output lại rồi mới in.

Kết quả là text hiện dần trên terminal:

Một Customer có thể đặt nhiều Order, còn mỗi Order chỉ thuộc về đúng một Customer...

Lưu ý: timeout(...) của HttpRequest chỉ áp dụng tới lúc nhận được response headers. Nó không giới hạn thời gian đọc body. Nếu model “treo” giữa chừng, vòng lặp sẽ chờ mãi. Đây là một lý do để chuyển sang WebClient trong ứng dụng thật.

5. Streaming với Spring WebClient

Trong một service Spring Boot, WebClient cho bạn một Flux<String> có hỗ trợ backpressure, timeout và cancel. Bạn cần thêm dependency:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-webflux</artifactId>
</dependency>

Cấu hình kết nối được đặt trong application.properties:

llm.base-url=${LLM_BASE_URL}
llm.api-key=${LLM_API_KEY}
llm.model=gpt-4o-mini

Tiếp theo, khai báo bean WebClient:

@Configuration
class LlmClientConfig {

    @Bean
    WebClient llmWebClient(@Value("${llm.base-url}") String baseUrl,
                           @Value("${llm.api-key}") String apiKey) {
        return WebClient.builder()
                .baseUrl(baseUrl)
                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .build();
    }
}

Service gọi API và trả về một Flux chứa các token:

@Service
public class ChatStreamService {

    private static final Logger log = LoggerFactory.getLogger(ChatStreamService.class);
    private static final ParameterizedTypeReference<ServerSentEvent<String>> SSE_TYPE =
            new ParameterizedTypeReference<>() {};

    private final WebClient llmWebClient;
    private final String model;
    private final ChatChunkParser parser = new ChatChunkParser();

    public ChatStreamService(WebClient llmWebClient, @Value("${llm.model}") String model) {
        this.llmWebClient = llmWebClient;
        this.model = model;
    }

    public Flux<String> stream(String question) {
        Map<String, Object> body = Map.of(
                "model", model,
                "stream", true,
                "messages", List.of(Map.of("role", "user", "content", question)));

        return llmWebClient.post()
                .uri("/v1/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(body)
                .retrieve()
                .bodyToFlux(SSE_TYPE)
                .mapNotNull(ServerSentEvent::data)
                .takeWhile(data -> !ChatChunkParser.DONE.equals(data))
                .mapNotNull(data -> parser.parsePayload(data).orElse(null))
                .timeout(Duration.ofSeconds(30))
                .doOnCancel(() -> log.info("Client đã hủy stream cho câu hỏi: {}", question));
    }
}

Từng bước trong chuỗi xử lý:

  • bodyToFlux(SSE_TYPE) để Spring tự tách sự kiện theo dòng trống và bỏ tiền tố data:. Bạn không phải tự xử lý buffer.
  • takeWhile hoàn tất Flux khi gặp [DONE].
  • timeout(Duration.ofSeconds(30)) áp dụng cho khoảng cách giữa hai phần tử liên tiếp, kể cả phần tử đầu tiên. Nếu 30 giây không có token mới, Flux báo lỗi TimeoutException. Đây là thứ HttpClient ở mục 4 còn thiếu.
  • Khi API trả 4xx/5xx, retrieve() ném WebClientResponseException trước khi có token nào.

5.1. Lưu toàn bộ câu trả lời sau khi stream xong

Bạn thường cần lưu câu trả lời đầy đủ để ghi lịch sử hoặc tính chi phí. StringBuilder phải được tạo riêng cho từng subscription. Vì vậy, bạn bọc chuỗi xử lý trong Flux.defer:

public Flux<String> streamAndRecord(String question) {
    return Flux.defer(() -> {
        StringBuilder full = new StringBuilder();
        return stream(question)
                .doOnNext(full::append)
                .doOnComplete(() -> log.info("Câu trả lời đầy đủ ({} ký tự): {}",
                        full.length(), full));
    });
}

Nếu StringBuilder được khai báo bên ngoài defer, hai request đồng thời sẽ ghi chung vào một buffer. Lỗi này khó phát hiện khi test thủ công vì bạn thường chỉ gửi một request mỗi lần.

6. Đẩy stream về trình duyệt

Controller trả về Flux<ServerSentEvent<...>> với produces = text/event-stream:

public record ChatChunk(String text) {}

@RestController
class ChatController {

    private final ChatStreamService chatStreamService;

    ChatController(ChatStreamService chatStreamService) {
        this.chatStreamService = chatStreamService;
    }

    @GetMapping(path = "/api/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    Flux<ServerSentEvent<ChatChunk>> stream(@RequestParam("q") String question) {
        return chatStreamService.streamAndRecord(question)
                .map(token -> ServerSentEvent.builder(new ChatChunk(token)).event("token").build())
                .concatWith(Mono.just(ServerSentEvent.builder(new ChatChunk("")).event("done").build()))
                .onErrorResume(ex -> Mono.just(ServerSentEvent
                        .builder(new ChatChunk("Mô hình đang bận, vui lòng thử lại."))
                        .event("error").build()));
    }
}

Bạn có thể thử endpoint bằng curl -N. Cờ -N tắt buffer của curl:

curl -N "http://localhost:8080/api/chat/stream?q=Order%20la%20gi"
event:token
data:{"text":"Order"}

event:token
data:{"text":" là"}

event:token
data:{"text":" đơn hàng"}

event:done
data:{"text":""}

Có ba quyết định thiết kế đáng chú ý:

  • Token được bọc trong JSON (ChatChunk), không gửi chuỗi thô. Theo đặc tả SSE, phía nhận sẽ bỏ một dấu cách ngay sau data:. Token " là" gửi dạng thô sẽ thành "là" ở trình duyệt, và các từ dính liền vào nhau. Token chứa \n cũng bị tách thành nhiều dòng data:. Bọc trong JSON giúp tránh cả hai lỗi.
  • Lỗi được gửi dưới dạng sự kiện error. Khi byte đầu tiên đã được gửi đi, status 200 đã nằm trong headers. Bạn không thể đổi sang 500 nữa. Phía client phải lắng nghe sự kiện error.
  • Sự kiện done được gửi rõ ràng. EventSource của trình duyệt tự kết nối lại khi stream đóng. Client cần sự kiện done để biết câu trả lời đã xong và tự gọi close().

Khi người dùng đóng tab, WebFlux cancel subscription. Tín hiệu cancel lan ngược lên WebClient, Reactor Netty đóng kết nối tới LLM API, và nhà cung cấp ngừng sinh token. Nhờ đó bạn không trả tiền cho phần câu trả lời không ai đọc.

7. Viết test

7.1. Test parser

class ChatChunkParserTest {

    private final ChatChunkParser parser = new ChatChunkParser();

    @Test
    void givenContentChunk_whenParseLine_thenReturnsToken() {
        String line = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\" chào\"}}]}";

        assertThat(parser.parseLine(line)).contains(" chào");
    }

    @Test
    void givenDoneOrRoleOnlyChunk_whenParseLine_thenReturnsEmpty() {
        assertThat(parser.parseLine("data: [DONE]")).isEmpty();
        assertThat(parser.parseLine(
                "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}")).isEmpty();
        assertThat(parser.parseLine(": keep-alive")).isEmpty();
    }

    @Test
    void givenBrokenJson_whenParseLine_thenThrows() {
        assertThatThrownBy(() -> parser.parseLine("data: {oops"))
                .isInstanceOf(IllegalStateException.class)
                .hasMessageContaining("Chunk không hợp lệ");
    }
}

Test đầu tiên xác nhận dấu cách ở đầu token được giữ nguyên. Đây chính là chi tiết dễ mất nhất trong cả luồng.

7.2. Test controller với WebTestClient

Bạn mock ChatStreamService để test không gọi API thật. @MockitoBean có từ Spring Boot 3.4. Với phiên bản cũ hơn, bạn dùng @MockBean.

@WebFluxTest(ChatController.class)
class ChatControllerTest {

    @Autowired
    WebTestClient webTestClient;

    @MockitoBean
    ChatStreamService chatStreamService;

    @Test
    void whenStreaming_thenSendsTokensThenDone() {
        when(chatStreamService.streamAndRecord("hi"))
                .thenReturn(Flux.just("Xin", " chào", "!"));

        List<ServerSentEvent<String>> events = webTestClient.get()
                .uri("/api/chat/stream?q=hi")
                .accept(MediaType.TEXT_EVENT_STREAM)
                .exchange()
                .expectStatus().isOk()
                .returnResult(new ParameterizedTypeReference<ServerSentEvent<String>>() {})
                .getResponseBody()
                .collectList()
                .block(Duration.ofSeconds(5));

        assertThat(events).extracting(ServerSentEvent::event)
                .containsExactly("token", "token", "token", "done");
        assertThat(events.get(1).data()).isEqualTo("{\"text\":\" chào\"}");
    }

    @Test
    void whenUpstreamFails_thenSendsErrorEvent() {
        when(chatStreamService.streamAndRecord("hi"))
                .thenReturn(Flux.concat(Flux.just("Xin"),
                        Flux.error(new TimeoutException())));

        List<ServerSentEvent<String>> events = webTestClient.get()
                .uri("/api/chat/stream?q=hi")
                .accept(MediaType.TEXT_EVENT_STREAM)
                .exchange()
                .returnResult(new ParameterizedTypeReference<ServerSentEvent<String>>() {})
                .getResponseBody()
                .collectList()
                .block(Duration.ofSeconds(5));

        assertThat(events).extracting(ServerSentEvent::event)
                .containsExactly("token", "error");
    }
}

Test thứ hai mô phỏng lỗi xảy ra giữa chừng. Client nhận được token đã sinh, sau đó là sự kiện error, và response vẫn mang status 200.

8. Lỗi hay gặp và khi nào không nên stream

  • Dùng bodyToMono(String.class). Cách này gom toàn bộ body lại rồi mới trả về, nên bạn mất hết lợi ích của streaming.
  • Reverse proxy buffer response. Nginx mặc định buffer response, và client sẽ nhận tất cả token cùng lúc ở cuối. Bạn cần tắt proxy_buffering cho route này, hoặc thêm header X-Accel-Buffering: no. Nén gzip ở tầng proxy cũng có thể gây buffer.
  • Dùng .block() trong controller Spring MVC. Nếu dự án dùng spring-boot-starter-web (Servlet), bạn vẫn có thể trả Flux hoặc dùng SseEmitter. Đừng block() rồi ghi từng phần, vì mỗi request sẽ giữ một thread trong suốt 30 giây.
  • Hiển thị output có cấu trúc khi chưa nhận đủ. Nếu bạn yêu cầu model trả JSON để parse hoặc gọi tool, các chunk chưa hoàn chỉnh không parse được. Lúc này bạn vẫn có thể stream, nhưng phải gom đủ rồi mới validate. Nếu không cần hiển thị dần, request thường (không stream) đơn giản hơn.
  • Tính token usage. Mặc định, chunk khi stream không có field usage. Nhiều API hỗ trợ tùy chọn kiểu "stream_options": {"include_usage": true} để gửi thêm một chunk cuối có usage và choices rỗng. Parser ở mục 3 bỏ qua chunk này an toàn vì path(0) không tồn tại.

Lưu ý: Không nên retry tự động khi đã gửi một phần token cho người dùng. Retry sẽ sinh một câu trả lời khác, nối vào phần cũ thành nội dung lộn xộn. Chỉ retry khi lỗi xảy ra trước token đầu tiên, ví dụ khi gặp 429 hoặc 503 từ retrieve().

9. Kết luận

Streaming response từ LLM API thực chất là đọc Server-Sent Events. Bạn đọc từng dòng data:, lấy choices[0].delta.content, và dừng khi gặp [DONE]. HttpClient của JDK đủ cho các công cụ nhỏ. Trong service Spring Boot, WebClient kết hợp Flux cho bạn thêm timeout giữa các token và khả năng cancel lan tới upstream.

Khi đẩy tiếp stream về trình duyệt, bạn nên bọc token trong JSON để giữ dấu cách và xuống dòng. Bạn cũng nên gửi các sự kiện done và error một cách rõ ràng, vì status code đã cố định ngay từ byte đầu tiên. Cuối cùng, hãy kiểm tra reverse proxy để chắc chắn nó không âm thầm buffer toàn bộ response.

llmspring-bootwebfluxserver-sent-eventsjava