package com.zcf.cn.chatspringboot.controller;

import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.BodyExtractors;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;

@RestController
public class ChatController {

    private final WebClient webClient;
    private static final Logger log = LoggerFactory.getLogger(ChatController.class); // 手动创建日志对象

    public ChatController(WebClient.Builder webClientBuilder) {
        // 配置全局默认请求头
        this.webClient = webClientBuilder
                .baseUrl("https://api.v0id.nl")
                .defaultHeader("Authorization", "Bearer hentai-vnsis8w8wsnsoskso") // 默认 Authorization
                .build();
    }

    @PostMapping(value = "/chat/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamChat(@RequestBody Map<String, Object> requestBody) {
        log.info("收到前端请求数据: {}", requestBody);
        HttpHeaders headers = new HttpHeaders();
        headers.set(HttpHeaders.CONNECTION, "keep-alive"); // 明确保持连接
        headers.set(HttpHeaders.CONTENT_TYPE, MediaType.TEXT_EVENT_STREAM_VALUE); // 设置流式返回格式

        return webClient.post()
                .uri("/v1/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(requestBody) // 设置请求体
                .exchangeToFlux(response -> {
                    if (response.statusCode().is2xxSuccessful()) {
                        // 从响应体中提取流式数据块
                        return response.body(BodyExtractors.toDataBuffers())
                                .flatMap(dataBuffer -> {
                                    try {
                                        // 读取 DataBuffer
                                        byte[] bytes = new byte[dataBuffer.readableByteCount()];
                                        dataBuffer.read(bytes);
                                        String chunk = new String(bytes, StandardCharsets.UTF_8);
                                        log.info("从外部 API 收到流式响应块: {}", chunk);
                                        return Mono.just(chunk);
                                    } catch (Exception e) {
                                        log.error("读取流数据块出错: {}", e.getMessage(), e);
                                        return Mono.empty();
                                    } finally {
                                        // 确保释放 DataBuffer
                                        DataBufferUtils.release(dataBuffer);
                                    }
                                });
                    } else {
                        return response.createException().flatMapMany(Mono::error);
                    }
                })
                .filter(chunk -> !chunk.trim().equals("data: [DONE]")) // 过滤结束标志
                .doOnComplete(() -> log.info("流式响应处理完成"))
                .doOnError(error -> log.error("处理流式响应时出错: {}", error.getMessage(), error))
                .doFinally(signalType -> log.info("流式请求结束，信号: {}", signalType));
    }

}



打印如下：
2024-12-13T15:34:45.107+08:00  INFO 26556 --- [0.0-8000-exec-1] c.z.c.c.controller.ChatController        : 收到前端请求数据: {model=gpt-4o-mini, stream=true, messages=[{role=user, content=写一首诗}]}
2024-12-13T15:34:49.392+08:00  INFO 26556 --- [ctor-http-nio-2] c.z.c.c.controller.ChatController        : 从外部 API 收到流式响应块: data: {"id":"chatcmpl-5b454124-6c6b-4428-8a65-4356ce577e23","object":"chat.completion.chunk","created":1734075288,"model":"gpt-4o-mini","choices":[{"delta":{"content":"在"},"index":0,"finish_reason":null}]}


2024-12-13T15:34:49.396+08:00  INFO 26556 --- [         task-1] c.z.c.c.controller.ChatController        : 流式请求结束，信号: cancel


为什么只有一块数据